/**
 *  TinyPasswordManager
 *  Copyright (C) 2008 Gross David <gdavid.devel@gmail.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "tinypasswordmanager.h"
#include "polarssl/aes.h"
#include "old_aes.h"
#include "sha256.h"

#include <fstream>
#include <cstring>
#include <algorithm>
#include <sstream>
#include <stdexcept>
#include <cctype>
#include <iostream>

#include <QtXml/QDomDocument>
#include <QtXml/QDomElement>

#ifdef WIN32
#define INITIALIZE_RAND() srand(GetTickCount());
#else
extern "C" {
#include <sys/time.h>
}

#define INITIALIZE_RAND() {\
struct timeval tv;\
gettimeofday(&tv, NULL);\
srand(tv.tv_usec);\
}
#endif

TinyPasswordManager::TinyPasswordManager()
    : m_cipheredBuff(NULL), m_filename("") {
    memset(m_key, 0, 32);
}

TinyPasswordManager::~TinyPasswordManager() {
    Identifier *tmp = NULL;

    while(!m_idents.empty()) {
        tmp = m_idents.back();
        m_idents.pop_back();
        delete tmp;
        tmp = NULL;
    }
    
    if(m_cipheredBuff) {
        delete[] m_cipheredBuff;
        m_cipheredBuff = NULL;
    }
}

bool TinyPasswordManager::isBeingUsed() const {
    return !m_filename.empty();
}

void TinyPasswordManager::createEmptyDatabase(const std::string &filename)
throw(ExBadFileFormat, ExFileErrorWriting) {
    if(isBeingUsed())
        reinitialize();

    m_filename = std::string(filename);
    save();
}

void TinyPasswordManager::openDatabase(const std::string &filename) 
throw(ExMemoryAllocationFailed, ExFileNotFound, ExFileErrorReading) {
    if(isBeingUsed() && !filename.compare(m_filename))
        return;
    else if(isBeingUsed() && filename.compare(m_filename))
        reinitialize();
        
    std::ifstream ifs(filename.c_str(), std::ios::in | std::ios::binary);

    if(!ifs.good())
        throw ExFileNotFound(__FUNCTION__, filename);

    ifs.seekg(0, std::ios::end);
    m_cipheredBuffSize = ifs.tellg();
    ifs.seekg(0, std::ios::beg);
    
    m_cipheredBuff = new unsigned char[m_cipheredBuffSize];
    if(!m_cipheredBuff)
        throw ExMemoryAllocationFailed(__FUNCTION__);
        
    ifs.read(reinterpret_cast<char *>(m_cipheredBuff), m_cipheredBuffSize);

    if(!ifs.good())
        throw ExFileErrorReading(__FUNCTION__, filename);

    m_filename = std::string(filename);
    ifs.close();
}

void TinyPasswordManager::getIdentifiers() 
throw(ExMemoryAllocationFailed, ExBadFileFormat) {
    if(!isBeingUsed())
        return;
            
    try {
        aes_context ctx;
        aes_setkey_dec(&ctx, m_key, 256);
        
        unsigned char iv[16];
        memset(iv, 0, 16);
        
        char *plaintext = new char[m_cipheredBuffSize];
        if(!plaintext)
            throw ExMemoryAllocationFailed(__FUNCTION__);
        
        if(aes_crypt_cbc(&ctx, AES_DECRYPT, m_cipheredBuffSize, iv, m_cipheredBuff, 
                         reinterpret_cast<unsigned char *>(plaintext)))
            throw ExBadFileFormat(__FUNCTION__, m_filename);

        unsigned int padding = plaintext[m_cipheredBuffSize - 1];
        
        if(padding > 16)
            throw ExBadFileFormat(__FUNCTION__, m_filename);

        plaintext[m_cipheredBuffSize - padding] = 0;

        QDomDocument doc;
        doc.setContent(QString(plaintext), false);
        QDomElement tpm = doc.documentElement();
        
        if(tpm.isNull())
            throw ExBadFileFormat(__FUNCTION__, m_filename);
            
        QDomElement element = tpm.firstChildElement();

        while(!element.isNull()) {
            if(element.tagName() == "identifier") {
                std::string login = element.firstChildElement("login").text().toStdString();
                std::string password = element.firstChildElement("password").text().toStdString();
                std::string description = element.firstChildElement("description").text().toStdString();
            
                Identifier *identifier = new Identifier(login, password, description);
                if(!identifier)
                    throw ExMemoryAllocationFailed(__FUNCTION__);
                    
                m_idents.push_back(identifier);
            }
            
            element = element.nextSiblingElement();
        }

        memset(plaintext, 0, m_cipheredBuffSize);
        delete[] plaintext;
        plaintext = NULL;
    }
    catch(ExBadFileFormat&) {
        AES aes;
        aes.SetParameters(256, 128);
        aes.StartDecryption(m_key);

        char *plaintext = new char[m_cipheredBuffSize];
        if(!plaintext)
            throw ExMemoryAllocationFailed(__FUNCTION__);
            
        aes.Decrypt(m_cipheredBuff, 
                    reinterpret_cast<unsigned char *>(plaintext), 
                    m_cipheredBuffSize >> 4, 
                    AES::CBC);
                    
        unsigned int padding = plaintext[m_cipheredBuffSize - 1];
        if(padding > 16)
            throw ExBadFileFormat(__FUNCTION__, m_filename);

        plaintext[m_cipheredBuffSize - padding] = 0;
        if(plaintext[0] != 'P' || plaintext[1] != 'M' || plaintext[2] != '1')
            throw ExBadFileFormat(__FUNCTION__, m_filename);

        char *tmp = strtok(plaintext + 16, "\n");
        if(!tmp)
            return;

        #define STAGE_LOGIN 0
        #define STAGE_PASSWORD 1
        #define STAGE_DESCRIPTION 2

        unsigned int stage = STAGE_LOGIN;
        Identifier *identifier = NULL;

        do {
            switch(stage % 3) {
            case STAGE_LOGIN:
                identifier = new Identifier();
                if(!identifier)
                    throw ExMemoryAllocationFailed(__FUNCTION__);
                    
                identifier->setLogin(tmp);
                break;

            case STAGE_PASSWORD:
                identifier->setPassword(tmp);
                break;

            case STAGE_DESCRIPTION:
                identifier->setDescription(tmp);
                m_idents.push_back(identifier);
                break;
            }

            stage++;
        } while((tmp = strtok(NULL, "\n")));

        memset(plaintext, 0, m_cipheredBuffSize);
        delete[] plaintext;
        plaintext = NULL;
    }

    std::sort(m_idents.begin(), m_idents.end(), TinyPasswordManager::compareIdentifier);
}

void TinyPasswordManager::save()
throw(ExMemoryAllocationFailed, ExBadFileFormat, ExFileErrorWriting) {
    if(!isBeingUsed())
        return;

    QDomDocument doc;
    
    QDomNode decl = doc.createProcessingInstruction("xml","version=\"1.0\" encoding=\"UTF-8\"");
    doc.insertBefore(decl, doc.firstChild());

    QDomElement tpm = doc.createElement("tpm");

    for(std::vector<Identifier*>::const_iterator it = m_idents.begin();
        it != m_idents.end(); ++it) {
        QDomElement identifier = doc.createElement("identifier");
        
        QDomElement elementLogin = doc.createElement("login");
        QString login = QString::fromStdString((*it)->getLogin());
        elementLogin.appendChild(doc.createTextNode(login));
        identifier.appendChild(elementLogin);
        
        QDomElement elementPassword = doc.createElement("password");
        QString password = QString::fromStdString((*it)->getPassword());
        elementPassword.appendChild(doc.createTextNode(password));
        identifier.appendChild(elementPassword);
        
        QDomElement elementDescription = doc.createElement("description");
        QString description = QString::fromStdString((*it)->getDescription());
        elementDescription.appendChild(doc.createTextNode(description));
        identifier.appendChild(elementDescription);
    
        tpm.appendChild(identifier);
    }

    doc.appendChild(tpm);
    std::string plaintext = doc.toString().toStdString();
        
    unsigned int padding = 16 - (plaintext.size() % 16);
    m_cipheredBuffSize = plaintext.size() + padding;

    char *decipheredBuff = new char[m_cipheredBuffSize];
    if(!decipheredBuff)
        throw ExMemoryAllocationFailed(__FUNCTION__);
                    
    strcpy(decipheredBuff, plaintext.c_str());
    addPadding(decipheredBuff, plaintext.size());      

    aes_context ctx;
    aes_setkey_enc(&ctx, m_key, 256);
    
    unsigned char iv[16];
    memset(iv, 0, 16);
    
    if(m_cipheredBuff)
        delete[] m_cipheredBuff;
        
    m_cipheredBuff = new unsigned char[m_cipheredBuffSize];
    if(!m_cipheredBuff)
        throw ExMemoryAllocationFailed(__FUNCTION__);
        
    if(aes_crypt_cbc(&ctx, AES_ENCRYPT, m_cipheredBuffSize, iv, reinterpret_cast<unsigned char *>(decipheredBuff), m_cipheredBuff))
        throw ExBadFileFormat(__FUNCTION__, m_filename);
        
    delete[] decipheredBuff;
    
    std::ofstream ofs(m_filename.c_str(), std::ios::out | std::ios::binary);

    ofs.seekp(std::ios_base::beg);
    ofs.write(reinterpret_cast<const char *>(m_cipheredBuff), m_cipheredBuffSize);

    if(!ofs.good())
        throw ExFileErrorWriting(__FUNCTION__, m_filename);

    ofs.close();
}

void TinyPasswordManager::addIdentifier(const Identifier &identifier) 
throw(ExMemoryAllocationFailed) { 
    Identifier *ident = new Identifier(identifier);
    if(!ident)
        throw ExMemoryAllocationFailed(__FUNCTION__);
        
    m_idents.push_back(ident);
    std::sort(m_idents.begin(), m_idents.end(), TinyPasswordManager::compareIdentifier);
}

void TinyPasswordManager::removeIdentifier(const Identifier &identifier) {
    for(std::vector<Identifier *>::iterator it = m_idents.begin();
        it != m_idents.end(); ++it) 
        if(**it == identifier) {
            Identifier *tmp = *it;
            m_idents.erase(it);
            
            delete tmp;
            tmp = NULL;
            break;
        }
}

void TinyPasswordManager::removeIdentifier(unsigned int index) 
throw(ExIndexOutOfBounds) {
    if(index >= m_idents.size())
        throw ExIndexOutOfBounds(__FUNCTION__, "m_idents");

    std::vector<Identifier *>::iterator it = m_idents.begin();
    it += index;

    Identifier *tmp = *it;
    m_idents.erase(it);
    
    delete tmp;
    tmp = NULL;
}

unsigned int TinyPasswordManager::size() const {
    return static_cast<unsigned int>(m_idents.size());
}

void TinyPasswordManager::setPrivateKey(const unsigned char key[32]) {
    memcpy(m_key, key, 32);
}

Identifier &TinyPasswordManager::at(unsigned int index)
throw(ExIndexOutOfBounds) {
    try {
        return *(m_idents.at(index));
    }
    catch(std::out_of_range &) {
        throw ExIndexOutOfBounds(__FUNCTION__, "m_idents");
    }
}

Identifier TinyPasswordManager::at(unsigned int index) const 
throw(ExIndexOutOfBounds) {
    try {
        return *(m_idents.at(index));
    }
    catch(std::out_of_range&) {
        throw ExIndexOutOfBounds(__FUNCTION__, "m_idents");
    }
}

std::string TinyPasswordManager::generatePassword(unsigned int minSize, unsigned int maxSize, const std::string &charset) {
    unsigned int passwordSize = (rand() % (maxSize - minSize)) + minSize;
    std::string password;
    unsigned int charsetIndex;

    INITIALIZE_RAND()

    for(unsigned int i = 0; i < passwordSize; ++i) {
        charsetIndex = rand() % charset.size();
        password += charset.at(charsetIndex);
    }

    return password;
}

bool TinyPasswordManager::compareIdentifier(Identifier *ident1, Identifier *ident2) {
    std::string descr1 = ident1->getDescription();
    std::string descr2 = ident2->getDescription();

    std::transform(descr1.begin(), descr1.end(), descr1.begin(), (int(*)(int))std::toupper);
    std::transform(descr2.begin(), descr2.end(), descr2.begin(), (int(*)(int))std::toupper);

    return (descr1.compare(descr2) > 0) ? false : true;
}

void TinyPasswordManager::addPadding(char *message, unsigned int length) {
    unsigned int i;
    unsigned int sizePadding = 16 - (length % 16);

    INITIALIZE_RAND()

    for(i = 0; i < sizePadding - 1; ++i)
        message[length + i] = static_cast<char>(rand());

    message[length + i] = static_cast<char>(sizePadding);
}

void TinyPasswordManager::reinitialize() {
    m_filename = "";

    while(!m_idents.empty()) {
        Identifier *tmp = m_idents.back();
        m_idents.pop_back();
        delete tmp;
        tmp = NULL;
    }
    
    delete[] m_cipheredBuff;
    m_cipheredBuff = NULL;
}
